"""
python p2_solution_pyswarm_particle.py
"""


import json
import os
import numpy as np
import pyswarms as ps

# Paths for input and output files
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/pyswarm_particle/'
os.makedirs(output_path, exist_ok=True)

# Load data
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)
with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)
with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# Data processing
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

categories = list(inventory_dict.keys())
num_categories = len(categories)
warehouses_list = list(warehouse_dict.keys())
num_warehouses = len(warehouses_list)

# Fitness function
def fitness_function(solution):
    costs = []
    
    for particle_solution in solution:  # Iterate over each particle's solution
        warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}
        total_cost = 0

        for category_index, warehouse_index in enumerate(particle_solution.astype(int)):
            warehouse_id = warehouses_list[warehouse_index]
            category_id = categories[category_index]
            inventory = inventory_dict[category_id]
            sales = sales_dict[category_id]

            # Update warehouse utilization
            warehouse_utilization[warehouse_id]['inventory'] += inventory
            warehouse_utilization[warehouse_id]['sales'] += sales

            # Add to total cost if this warehouse is used
            if inventory > 0:
                total_cost += warehouse_dict[warehouse_id]['daily_cost']

        # Check constraints and apply penalties
        feasible = True
        for warehouse_id, usage in warehouse_utilization.items():
            max_inventory = warehouse_dict[warehouse_id]['max_inventory']
            max_sales = warehouse_dict[warehouse_id]['max_sales']
            if usage['inventory'] > max_inventory or usage['sales'] > max_sales:
                feasible = False
                break

        # Penalize if constraints are not met
        if feasible:
            costs.append(total_cost)
        else:
            costs.append(float('inf'))

    return np.array(costs)


# Run PSO
def run_pso():
    # Set bounds to be arrays of shape (num_categories,), where each entry is [0, num_warehouses - 1]
    lower_bounds = np.zeros(num_categories)
    upper_bounds = np.ones(num_categories) * (num_warehouses - 1)
    bounds = (lower_bounds, upper_bounds)
    
    # Initialize the optimizer with particle swarm parameters
    optimizer = ps.single.GlobalBestPSO(
        n_particles=50,
        dimensions=num_categories,
        options={'c1': 0.5, 'c2': 0.3, 'w': 0.9},
        bounds=bounds
    )

    # Run optimization
    best_cost, best_pos = optimizer.optimize(fitness_function, iters=100)
    return best_cost, best_pos.astype(int)


# Execute PSO and extract the best solution
best_cost, best_solution = run_pso()

# Calculate warehouse utilization and save results
warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}
assignment = []

for category_index, warehouse_index in enumerate(best_solution):
    warehouse_id = warehouses_list[warehouse_index]
    category_id = categories[category_index]
    inventory = inventory_dict[category_id]
    sales = sales_dict[category_id]
    
    # Update utilization for the assigned warehouse
    warehouse_utilization[warehouse_id]['inventory'] += inventory
    warehouse_utilization[warehouse_id]['sales'] += sales
    
    # Store assignment result
    assignment.append({'category_id': category_id, 'warehouse_id': warehouse_id})

# Save assignments
with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
    json.dump(assignment, f)

# Save total cost
with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
    f.write(str(best_cost))

# Calculate and save utilization rates
utilization_rates_inventory = []
utilization_rates_sales = []

for warehouse_id, usage in warehouse_utilization.items():
    max_inventory = warehouse_dict[warehouse_id]['max_inventory']
    max_sales = warehouse_dict[warehouse_id]['max_sales']
    
    utilization_rates_inventory.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_inventory': usage['inventory'] / max_inventory
    })
    utilization_rates_sales.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_sales': usage['sales'] / max_sales
    })

with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
    json.dump(utilization_rates_inventory, f)

with open(os.path.join(output_path, 'all_utilization_rate_sales.json'), 'w') as f:
    json.dump(utilization_rates_sales, f)

print("Optimization Complete. Best total cost:", best_cost)
